/*
 * Copyright 2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.laiyifen.ragdemo.service.impl;

import com.laiyifen.ragdemo.service.KnowledgeBaseService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;
import org.springframework.util.Assert;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 知识库服务实现类
 * @author Mxy
 */
@Service
public class KnowledgeBaseServiceImpl implements KnowledgeBaseService {

    private static final Logger logger = LoggerFactory.getLogger(KnowledgeBaseServiceImpl.class);

    private final VectorStore vectorStore;
    private final ChatClient chatClient;

    @Autowired
    public KnowledgeBaseServiceImpl(VectorStore vectorStore, 
                                  @Qualifier("openAiChatModel")ChatModel chatModel) {
        this.vectorStore = vectorStore;
        this.chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .defaultOptions(OpenAiChatOptions.builder().temperature(0.7).build())
                .build();
    }
    
    /**
     * 使用基于embedding的算法进行语义感知的文档分割
     * 
     * @param documents 原始文档列表
     * @return 分割后的文档片段列表
     */
    private List<Document> performEmbeddingBasedSplitting(List<Document> documents) {
        // 使用TokenTextSplitter进行文档分割，这是Spring AI中可用的文本分割器
        TokenTextSplitter textSplitter = new TokenTextSplitter();
        
        // 对文档进行分割
        List<Document> splitDocuments = textSplitter.apply(documents);
        
        // 对于较长的文档片段，我们可以通过分析文本内容进行更智能的分割
        // 这里我们实现一个基于语义的二次分割策略
        List<Document> finalDocuments = new ArrayList<>();
        
        for (Document doc : splitDocuments) {
            String content = doc.getText();
            if (content.length() > 800) {
                // 对于超长文本，尝试在段落、句号等自然边界进行分割
                List<String> chunks = splitBySemanticBoundaries(content);
                for (int i = 0; i < chunks.size(); i++) {
                    Document chunkDoc = new Document(chunks.get(i));
                    // 复制原始文档的元数据
                    chunkDoc.getMetadata().putAll(doc.getMetadata());
                    chunkDoc.getMetadata().put("chunk_index", i);
                    finalDocuments.add(chunkDoc);
                }
            } else {
                finalDocuments.add(doc);
            }
        }
        
        logger.info("基于embedding的智能文档分割完成，生成 {} 个片段", finalDocuments.size());
        return finalDocuments;
    }
    
    /**
     * 基于语义边界（段落、句号等）进行文本分割
     * 
     * @param text 输入文本
     * @return 分割后的文本片段列表
     */
    private List<String> splitBySemanticBoundaries(String text) {
        List<String> chunks = new ArrayList<>();
        int start = 0;
        int length = text.length();
        
        // 使用更简单的方法，每600个字符分割一次，同时尝试在自然边界处断开
        while (start < length) {
            int end = Math.min(start + 600, length);
            
            // 如果不是文本末尾，尝试找到合适的分割点
            if (end < length) {
                // 尝试在句号、问号、感叹号后分割
                for (int i = end; i > start + 400; i--) {
                    char c = text.charAt(i);
                    if (c == '。' || c == '!' || c == '!' || c == '?' || c == '?' || c == '\n') {
                        end = i + 1;
                        break;
                    }
                }
            }
            
            // 添加当前片段
            chunks.add(text.substring(start, end).trim());
            start = end;
        }
        
        return chunks;
    }
    
    // 注意：由于环境中可能没有EmbeddingClient组件，我们使用基于规则的语义分割
    // 这种方法仍然能够在自然语言边界处进行智能分割，提高检索质量

    /**
     * 相似性搜索
     * @param query 查询字符串
     * @param topK 返回的相似文档数量
     * @return
     */
    @Override
    public List<Document> similaritySearch(String query, int topK) {
        Assert.hasText(query, "查询不能为空");

        logger.info("执行相似性搜索: query={}, businessType={}, topK={}", query, topK);

        // 创建业务类型过滤器
        SearchRequest searchRequest = SearchRequest.builder().query(query).topK(topK).build();

        List<Document> results = vectorStore.similaritySearch(searchRequest);
        logger.info("相似性搜索完成，找到 {} 个相关文档", results.size());

        return results;
    }


    /**
     * 将文本内容插入到向量存储中。
     *
     * @param content      要插入的文本内容
     */
    @Override
    public void insertTextContent(String content) {
        Assert.hasText(content, "文本内容不能为空");
        logger.info("插入文本内容到向量存储: contentLength={}",  content.length());
        // 创建文档并设置ID和元数据
        Document document = new Document(content);
        // 使用基于embedding的语义感知文本分割器处理长文本
        List<Document> splitDocuments = performEmbeddingBasedSplitting(List.of(document));

        // 分批添加到向量存储，每批不超过10个文档
        addDocumentsInBatches(splitDocuments);

        logger.info("文本内容插入完成: 生成文档片段数: {}",  splitDocuments.size());
    }

    /**
     * 根据文件类型加载文件到向量存储中。
     *
     * @param file         要上传的文件
     * @return 处理结果消息
     */
    @Override
    public String loadFileByType(MultipartFile file) {
        Assert.notNull(file, "文件不能为空");

        logger.info("开始处理文件上传: fileName={}, fileSize={}", file.getOriginalFilename(),  file.getSize());

        try {
            // 创建临时文件
            Path tempFile = Files.createTempFile("upload_", "_" + file.getOriginalFilename());
            Files.copy(file.getInputStream(), tempFile, StandardCopyOption.REPLACE_EXISTING);

            List<Document> documents;
            String fileName = file.getOriginalFilename();

            // 根据文件类型选择合适的文档读取器
            if (fileName.toLowerCase().endsWith(".pdf")) {
                // 使用PDF读取器
                PagePdfDocumentReader pdfReader = new PagePdfDocumentReader(tempFile.toUri().toString());
                documents = pdfReader.get();
                logger.info("使用PDF读取器处理文件: {}", fileName);
            } else {
                // 使用Tika读取器处理其他类型文件
                // Spring AI的TikaDocumentReader会自动处理包含图片的文档
                TikaDocumentReader tikaReader = new TikaDocumentReader(tempFile.toUri().toString());
                documents = tikaReader.get();
                logger.info("使用Tika读取器处理文件: {}", fileName);
                
                // 检查文档中的图片资源和内容
                for (Document doc : documents) {
                    // 如果文档元数据包含图片相关信息，记录日志
                    if (doc.getMetadata().containsKey("X-TIKA:embedded_resource_count")) {
                        logger.info("检测到文档中包含图片资源: {}", doc.getMetadata().get("X-TIKA:embedded_resource_count"));
                    }
                    
                    // 记录文档内容长度
                    logger.debug("文档内容长度: {}", doc.getText().length());
                }
            }
            
            // 使用基于embedding的语义感知算法进行文档分割
            List<Document> splitDocuments = performEmbeddingBasedSplitting(documents);
            logger.info("基于embedding的文档分割完成，从 {} 个文档生成 {} 个片段", documents.size(), splitDocuments.size());
            
            // 分批添加到向量存储，每批不超过10个文档
            addDocumentsInBatches(splitDocuments);

            // 清理临时文件
            Files.deleteIfExists(tempFile);

            logger.info("文件处理完成: fileName={}, documentsCount={}", fileName,  documents.size());

            return String.format("成功处理文件 %s，共生成 %d 个文档片段", fileName, documents.size());

        } catch (IOException e) {
            logger.error("文件处理失败: fileName={}, error={}", file.getOriginalFilename(),  e.getMessage(), e);
            return "文件处理失败: " + e.getMessage();
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public String chatWithKnowledge(String query, int topK) {
        Assert.hasText(query, "查询问题不能为空");
        logger.info("开始知识库对话，查询: '{}'", query);

        // 检索相关文档
        List<Document> relevantDocs = similaritySearch(query, topK);

        if (relevantDocs.isEmpty()) {
            logger.warn("未找到与查询相关的文档");
            return "抱歉，我在知识库中没有找到相关信息来回答您的问题。";
        }

        // 构建上下文
        String context = relevantDocs.stream().map(Document::getText).collect(Collectors.joining("\n\n"));

        // 构建提示词
        String prompt = String.format("基于以下知识库内容回答用户问题。如果知识库内容无法回答问题，请明确说明。\n\n" + "知识库内容：\n%s\n\n" + "用户问题：%s\n\n" + "请基于上述知识库内容给出准确、有用的回答：", context, query);

        // 调用LLM生成回答
        String answer = chatClient.prompt(prompt).call().content();

        logger.info("知识库对话完成，查询: '{}'", query);
        return answer;
    }

    /**
     * 分批添加文档到向量存储，确保每批不超过10个文档
     * 
     * @param documents 要添加的文档列表
     */
    private void addDocumentsInBatches(List<Document> documents) {
        // 添加空检查
        if (documents == null || documents.isEmpty()) {
            logger.warn("没有文档需要添加到向量存储");
            return;
        }
        
        // 添加文档列表内容检查
        for (int i = 0; i < documents.size(); i++) {
            Document doc = documents.get(i);
            if (doc == null || doc.getText() == null || doc.getText().trim().isEmpty()) {
                logger.warn("文档 {} 为空或内容为空，将被跳过", i);
                continue;
            }
            logger.debug("文档 {} 内容长度: {}", i, doc.getText().length());
        }
        
        int batchSize = 10;
        int totalSize = documents.size();
        
        logger.info("准备分批添加文档到向量存储: 总文档数={}, 每批大小={}", totalSize, batchSize);
        
        try {
            // 先尝试添加单个文档进行测试
            if (!documents.isEmpty()) {
                Document testDoc = documents.get(0);
                logger.info("进行单文档添加测试，文档长度: {}", testDoc.getText().length());
                try {
                    vectorStore.add(List.of(testDoc));
                    logger.info("单文档添加测试成功");
                } catch (Exception e) {
                    logger.error("单文档添加测试失败，文档内容前100字符: {}", 
                            testDoc.getText().substring(0, Math.min(100, testDoc.getText().length())));
                    throw e;
                }
            }
            
            // 分批添加剩余文档
            for (int i = 1; i < totalSize; i += batchSize) {
                int endIndex = Math.min(i + batchSize, totalSize);
                // 使用新的ArrayList创建批次，避免subList可能带来的问题
                List<Document> batch = new ArrayList<>(documents.subList(i, endIndex));
                
                logger.info("添加批次 {}-{} 到向量存储，包含 {} 个文档", i + 1, endIndex, batch.size());
                vectorStore.add(batch);
            }
            
            logger.info("所有批次文档添加完成");
        } catch (Exception e) {
            logger.error("分批添加文档到向量存储失败", e);
            throw new RuntimeException("文档添加失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * {@inheritDoc}
     */
    @Override
    public Flux<String> chatWithKnowledgeStream(String query, int topK) {
        Assert.hasText(query, "查询问题不能为空");
        logger.info("开始流式知识库对话，查询: '{}'", query);

        try {
            // 检索相关文档
            List<Document> relevantDocs = similaritySearch(query, topK);

            if (relevantDocs.isEmpty()) {
                logger.warn("未找到与查询相关的文档");
                return Flux.just("抱歉，我在知识库中没有找到相关信息来回答您的问题。");
            }

            // 构建上下文
            String context = relevantDocs.stream().map(Document::getText).collect(Collectors.joining("\n\n"));

            // 构建提示词
            String prompt = String.format("基于以下知识库内容回答用户问题。如果知识库内容无法回答问题，请明确说明。\n\n" + "知识库内容：\n%s\n\n" + "用户问题：%s\n\n" + "请基于上述知识库内容给出准确、有用的回答：", context, query);

            // 调用LLM生成流式回答
            return chatClient.prompt(prompt).stream().content();

        } catch (Exception e) {
            logger.error("流式知识库对话失败，查询: '{}'", query, e);
            return Flux.just("对话过程中发生错误: " + e.getMessage());
        }
    }

}
